import time
import socket
import pynvml
import redis
from collections import deque
import setting
import os
from datetime import datetime

MONITOR_REDIS = os.environ.get('MONITOR_REDIS')
MONITOR_REDIS_PORT = int(os.environ.get('MONITOR_REDIS_PORT'))
# 初始化NVIDIA管理库
pynvml.nvmlInit()

# 获取本机IP地址
def get_host_ip():
    hostname = socket.gethostname()
    ip_address = socket.gethostbyname(hostname)
    return ip_address

# 获取本机名称
def get_host_name():
    return socket.gethostname()

# 获取GPU信息
def get_gpu_info():
    # 假设我们想要监控所有GPU
    num_gpus = pynvml.nvmlDeviceGetCount()
    host_info_list = []
    gpu_info_list = []
    for i in range(num_gpus):
        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
        gpu_model = pynvml.nvmlDeviceGetName(handle)
        gpu_memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
        gpu_memory_total = gpu_memory_info.total / (1024 ** 2)  # 转换为MB
        gpu_memory_used = gpu_memory_info.used / (1024 ** 2)  # 转换为MB
        current_time = time.strftime("%H:%M:%S", time.localtime())

        gpu_info = {
            'time': current_time,
            'model': gpu_model,
            'memory_total': gpu_memory_total,
            'memory_used': gpu_memory_used,
        }
        gpu_info_list.append(gpu_info)
    return gpu_info_list

# 初始化Redis连接
redis_client = redis.Redis(host=MONITOR_REDIS, port=MONITOR_REDIS_PORT, db=0)

# 监控函数
def monitor_gpus():
    
    host_name = get_host_name()
    ip = get_host_ip()
    gpus = [deque(maxlen=60)]
    saveNum = 0
    while True:
        gpu_infos = get_gpu_info()
        for i in range(len(gpu_infos)):
            gpu_info = gpu_infos[i]
            # 将数据写入Redis并设置60秒过期时间
            gpus[i].append(gpu_info)
        time.sleep(1)  # 等待1秒后再进行下一次监控
        saveNum += 1
        # 每5秒保存一次
        if(saveNum % 5 == 0):
            # 把每一次监控的数据都写入到监控里机
            # 使用主机名和GPU型号作为Redis键的一部分
            for j in range(len(gpus)):
                ajson = gpus[j][0]
                model = gpus[j][0]["model"]
                redis_key = f"gpu_monitor:{host_name}_{model}"
                # 将内存总量和已用内存作为Redis值
                redis_value = f"{list(gpus[j])}"
                redis_client.setex(redis_key, 1200, redis_value)
        saveNum = 0 if(saveNum == 60) else saveNum
# 开始监控
if __name__ == "__main__":
    monitor_gpus()